"""
Automaton module. Supports easy state-machines usage.

'automaton' decorator adds following properties to automaton class:
    event_list      - list of provided events
    state_list      - list of available states
    state           - current state, None if automaton is finished
    isFinished()    - same as condition (state == None)
    restart()       - restart finished automaton
    caseEvent(e)    - send event to automaton

'event' constructor provides multiple interfaces
    1. Simple events with or without parameters:
        x = event()
        y = event('message, info')
        z = event('x, y, z', kw_props = 'default-value')

    2. User-defined initialization:
        @event
        def my_event(self, x, y):
            self.x = x
            self.y = y
            self.sum = x + y

    3. User-defined class:
        class my_event(Event):
            def __init__(self, x, y):
                pass

'initial' decorator marks class that it should be initial. Automaton should have exectly one initial state

'on_event' decorator provide generic event handler interface:
    @on_event('an_event')
    @on_event('an_event', next = 'NEW_STATE')
    @on_event('an_event', next = ['NEW_STATE', 'NEXT_STATE_2']) # transition depends on user logic
    @on_event('command', lambda a, e: a.counter >= 10 and e.cmd == 'tick', 'NEW_STATE')

'default' default event handler decorator:
    @default
    @defalut(lambda a, e: a.isAccept(e))
    @defalut(next = "NEXT_STATE")

on_XXX and default handlers with annotations. XXX used as event name, annotation is parsed as transition and additional condition:
    on_event(self, e)           # same as @on_event('event')
    on_select(self, e : lambda a, e: e.selection is not None) -> 'NEXT_STATE'
                                # same as @on_event('select', lambda a, e: ..., 'NEXT_STATE')
    default(self, e, lambda a, e: ...) -> ['STATE_1', 'STATE_2']

Example:
    @automaton
    class Switcher:
        on = event()
        off = event()
        get_value = event()

        @initial
        class OFF:
            def onEntry(self):
                self.status = "Switcher is off."

            def on_on(self, e) -> "ON":
                pass

        class ON:
            def onEntry(self):
                self.status = "Switch is on."

            def on_off(self, e) -> "OFF":
                pass

            def on_get_value(self, e):
                return 10

Usage:
    switcher = Switcher()
    assert switcher.state == Switcher.OFF
    assert switcher.status == "Switcher is off."

    switcher.on()
    assert switcher.state == Switcher.ON
    assert switcher.status == "Switcher is on."

    assert switcher.getvalue() == 10
"""
from S3D.Decorators import *
from inspect import isfunction

ENABLE_CONSTR = 0

# TODO: event handlers are too ugly
# TODO: move tests to separate file
# TODO: on_event makes event handler
# TODO: take handler parameters from annotations

class Event:
    pass

ctr_event = ctr_type(Event)

class EventWrapper:
    __slots__ = ('event')

    @constraints(ENABLE_CONSTR)
    def __init__(self, event : ctr_event):
        self.event = event

    def __call__(self, *args, **kw):
        return self.event(*args, **kw)

    def __get__(self, inst, _):
        if inst is not None:
            return EventBind(inst, self.event)
        return self.event

class EventBind:
    __slots__ = ('auto', 'event')

    @constraints(ENABLE_CONSTR)
    def __init__(self, auto : ctr_has('castEvent'), event : type):
        self.auto   = auto
        self.event  = event

    def __call__(self, *args, **kw):
        e = self.event(*args, **kw)
        return self.auto.castEvent(e)

@constraints(ENABLE_CONSTR)
def event(decl : [str, isfunction] = "", **kw) -> EventWrapper:
    if isfunction(decl):
        assert len(kw) == 0

        class UserEvent(Event):
            pass

        UserEvent.__name__      = decl.__name__
        UserEvent.__module__    = decl.__module__
        UserEvent.__init__      = decl
        return EventWrapper(UserEvent)

    decl_list = tuple(decl.replace(',', ' ').split())
    assert len(decl_list) == len(set(decl_list)), 'Duplication in event arguments'
    assert len(set(decl_list) | set(kw)) == len(decl_list) + len(kw), 'Intersection between args and keywords'

    class UserEvent(Event):
        __slots__ = decl_list + tuple(kw)

        def __init__(self, *args, **newKw):
            if len(args) > len(decl_list):
                raise AssertionError('Too many arguments: %s' % args[len(decl_list):])
            if len(args) + len(newKw) < len(decl_list):
                raise AssertionError('Arguments expected: %s' % decl_list[len(args):])
            if len(args) + len(newKw) > len(decl_list) + len(kw):
                raise AssertionError('Unexpected arguments: %s' % [p for p in newKw if p not in kw])

            val_list = list(args)
            for index in range(len(val_list), len(decl_list)):
                val_list.append(newKw[decl_list[index]])

            for name, val in zip(decl_list, val_list):
                setattr(self, name, val)

            for name, val in kw.items():
                setattr(self, name, newKw.get(name) or val)

    return EventWrapper(UserEvent)

class EventHandler:
    __slots__ = ('event', 'cond', 'handler', 'index', 'next')

    counter = 0

    def __init__(self, event, cond, handler, next):
        EventHandler.counter += 1
        self.index   = EventHandler.counter
        self.event   = event
        self.cond    = cond
        self.handler = handler
        self.next    = next

    def __lt__(self, e):
        if self.event is None and e.event is not None:
            return False
        elif self.event is not None and e.event is None:
            return True
        return self.index < e.index

    def initialize(self, cls):
        if isinstance(self.event, str):
            self.event = getattr(cls, self.event)

        # make checker
        if isinstance(self.next, list):
            assert isinstance(self.next, list), self.next
            self.next = [ getattr(cls, p) if isinstance(p, str) else p for p in self.next ]

            next_list = self.next
            if len(next_list) == 0:
                self.next = lambda prev, next: next is prev
            elif len(next_list) == 1:
                next_event = next_list[0]
                self.next = lambda prev, next: next is prev or next is next_event
            else:
                self.next = lambda prev, next: prev is next or any(map(lambda x: next is x, next_list))

@constraints(ENABLE_CONSTR)
def on_event(
    event : [str, ctr_event],
    cond_or_next : [isfunction, str, [str]] = None,
    next : [str, [str]] = None
):
    if isfunction(cond_or_next):
        cond = cond_or_next
    else:
        assert next is None
        cond = None
        next = cond_or_next

    if next is None:
        next = []

    def decor(func):
        return EventHandler(event, cond, func, next)
    return decor

@constraints(ENABLE_CONSTR)
def isHandlerName(func : ctr_call()) -> bool:
    return func.__name__.startswith('on_') or func.__name__ == 'default'

@constraints(ENABLE_CONSTR)
def event_handler(func : ctr_call()) -> EventHandler:
    if not isHandlerName(func):
        raise AssertionError('Invalid function name: %s' % func.__name__)

    eventName = func.__name__[3:] if func.__name__ != 'default' else None

    cond = None
    for k, v in func.__annotations__.items():
        if k != 'return' and isfunction(v):
            cond = v
            break

    next = func.__annotations__['return'] if 'return' in func.__annotations__ else []
    if not constr([None, str, [str]])(next):
        raise AssertionError('Invalid next state: %s' % next)

    return EventHandler(eventName, cond, func, next)

@decorator
@constraints(ENABLE_CONSTR)
def default(
    func, *,
    cond : ctr_call() = None,
    next : [str, [str]] = []
):
    return EventHandler(None, cond, func, next)

class InitialStateWrapper:
    __slots__ = ('state')

    def __init__(self, state):
        self.state = state

def initial(cls):
    return InitialStateWrapper(cls)

class AutomatonState:
    def __get__(self, inst, _):
        assert inst is not None, 'Getting state from automaton class'
        return inst.__state__

    def __set__(self, inst, value):
        assert inst.isCastingEvent

        if value is None:
            inst.__state__ = None
            return

        inst.__state__ = getattr(inst, value)

_AUTOMATON_PROPS = [
    'event_list',
    'state_list',
    'handler_list',
    'state',
    'restart',
    'finish'
    'isFinished',
    'isCastingEvent',
    'castEvent',
]

@constraints(ENABLE_CONSTR)
def automaton(cls : type) -> type:
    for prop in _AUTOMATON_PROPS:
        assert not hasattr(cls, prop), 'Automaton property conflict: %s' % prop

    # make sure it doesn't have automatons properties
    event_list = []
    state_list = []
    initialState = None
    handler_list = []

    #
    # Process states and events
    #
    for name in dir(cls):
        if name.startswith('_'):
            continue

        val = getattr(cls, name)

        if isinstance(val, InitialStateWrapper):
            setattr(cls, name, val.state) # remove wrapper
            assert initialState is None, 'Too many initial states'
            initialState = name

        if isfunction(val) and isHandlerName(val):
            val = event_handler(val)
            setattr(cls, name, val)

        if isinstance(val, EventHandler):
            val.initialize(cls)
            handler_list.append(val)
            continue

        if not isinstance(val, type):
            continue

        if issubclass(val, Event):
            val.__name__ = name # make sure that an event has correct name
            setattr(cls, name, EventWrapper(val)) # make sure wrapper is used for every event
            event_list.append(name)
            continue

        state_list.append(name)

    handler_list.sort()
    cls.handler_list = handler_list

    # initial state
    if len(state_list) == 1 and initialState is None:
        initialState = state_list[0]
    elif len(state_list) > 0 or initialState is not None:
        assert initialState is not None, 'Initial state not found'
        state_list.insert(0, initialState)

    #
    # Process handlers
    #
    for stateName in state_list:
        state = getattr(cls, stateName)

        assert not hasattr(state, 'handler_list')

        handler_list = []
        for name in dir(state):
            if name.startswith('_'):
                continue

            val = getattr(state, name)

            if isfunction(val) and isHandlerName(val):
                val = event_handler(val)
                setattr(state, name, val)
            elif not isinstance(val, EventHandler):
                continue

            val.initialize(cls)

            handler_list.append(val)

        handler_list.sort()
        state.handler_list = tuple(handler_list)

    cls.event_list = tuple(event_list)
    cls.state_list = tuple(state_list)

    cls.__init_original__ = cls.__init__

    @method_of(cls)
    def __init__(self, *args, **kw):
        self.__init_original__(*args, **kw)
        self.isCastingEvent = True
        self.state = None
        self.isCastingEvent = False
        self.restart()

    @method_of(cls)
    def isFinished(self):
        return self.state is None

    @method_of(cls)
    def restart(self):
        assert self.isFinished()

        self.onInitial()

        if len(self.state_list) == 0:
            self.onFinish()
            return
        
        self.isCastingEvent = True
        self.state = self.state_list[0]
        self.isCastingEvent = False

        self.state.onEnter(self)

    @method_of(cls)
    def finish(self):
        assert self.isCastingEvent
        self.state = None

    def buildTransitionName(prev, next, h):
        if h.event is None:
            return '%s -> %s' % (prev.__name__, next.__name__)
        return '%s(%s) -> %s' % (prev.__name__, h.event.__name__, next.__name__)

    def concat(*args):
        for p in args:
            for it in p:
                yield it


    @method_of(cls)
    def castEvent(self, e):
        if self.isFinished():
            return None

        self.isCastingEvent = True

        prev = self.state

        for h in concat(self.state.handler_list, self.handler_list):
            if h.event is not None and not isinstance(e, h.event):
                continue
            if h.cond is not None and not h.cond(self, e):
                continue

            r = h.handler(self, e)

            next = self.state
            if next is not None:
                if h.next is None or isinstance(h.next, str):
                    if next is not prev:
                        self = prev.__name__ # restore previous state
                        raise AssertionError('Transition is triggered automatically: %s' % buildTransitionName(prev, next, h))
                    self.state = h.next
                else:
                    if not h.next(prev, self.state):
                        self = prev.__name__ # restore previous state
                        raise AssertionError('Invalid transition: %s' % buildTransitionName(prev, next, h))

            if self.state != prev:
                prev.onLeave(self)

                if self.state is not None:
                    self.state.onEnter(self)
                else:
                    self.onFinish()

            self.isCastingEvent = False
            return r

        self.isCastingEvent = False
        return None

    cls.state = AutomatonState()

    for name in cls.state_list:
        state = getattr(cls, name)
        if not hasattr(state, 'onEnter'):
            state.onEnter = lambda x: None
        if not hasattr(state, 'onLeave'):
            state.onLeave = lambda x: None

    if not hasattr(cls, 'onInitial'):
        cls.onInitial = lambda x: None
    if not hasattr(cls, 'onFinish'):
        cls.onFinish = lambda x: None

    return cls
